import faiss
import numpy as np
import os
import sys

current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
print(f"current_dir: {current_dir}")
sys.path.insert(0, parent_dir)
from my_common import  load_flag_model, load_flag_reranker_large_model, get_openai_client

# 1. Data
corpus = [
    "Cheli: A downtown Chinese restaurant presents a distinctive dining experience with authentic and sophisticated flavors of Shanghai cuisine. Avg cost: $40-50",
    "Masa: Midtown Japanese restaurant with exquisite sushi and omakase experiences crafted by renowned chef Masayoshi Takayama. The restaurant offers a luxurious dining atmosphere with a focus on the freshest ingredients and exceptional culinary artistry. Avg cost: $500-600",
    "Per Se: A midtown restaurant features daily nine-course tasting menu and a nine-course vegetable tasting menu using classic French technique and the finest quality ingredients available. Avg cost: $300-400",
    "Ortomare: A casual, earthy Italian restaurant locates uptown, offering wood-fired pizza, delicious pasta, wine & spirits & outdoor seating. Avg cost: $30-50",
    "Banh: Relaxed, narrow restaurant in uptown, offering Vietnamese cuisine & sandwiches, famous for its pho and Vietnam sandwich. Avg cost: $20-30",
    "Living Thai: An uptown typical Thai cuisine with different kinds of curry, Tom Yum, fried rice, Thai ice tea, etc. Avg cost: $20-30",
    "Chick-fil-A: A Fast food restaurant with great chicken sandwich, fried chicken, fries, and salad, which can be found everywhere in New York. Avg cost: 10-20",
    "Joe's Pizza: Most famous New York pizza locates midtown, serving different flavors including classic pepperoni, cheese, spinach, and also innovative pizza. Avg cost: $15-25",
    "Red Lobster: In midtown, Red Lobster is a lively chain restaurant serving American seafood standards amid New England-themed decor, with fair price lobsters, shrips and crabs. Avg cost: $30-50",
    "Bourbon Steak: It accomplishes all the traditions expected from a steakhouse, offering the finest cuts of premium beef and seafood complimented by wine and spirits program. Avg cost: $100-150",
    "Da Long Yi: Locates in downtown, Da Long Yi is a Chinese Szechuan spicy hotpot restaurant that serves good quality meats. Avg cost: $30-50",
    "Mitr Thai: An exquisite midtown Thai restaurant with traditional dishes as well as creative dishes, with a wonderful bar serving cocktails. Avg cost: $40-60",
    "Yichiran Ramen: Famous Japenese ramen restaurant in both midtown and downtown, serving ramen that can be designed by customers themselves. Avg cost: $20-40",
    "BCD Tofu House: Located in midtown, it's famous for its comforting and flavorful soondubu jjigae (soft tofu stew) and a variety of authentic Korean dishes. Avg cost: $30-50",
]

user_input = "I want some Chinese food"
if __name__ == "__main__":
    # 2. Indexing
    """
    Now we need to figure out a fast but powerful enough method to retrieve docs in the corpus that are most closely related to our questions. Indexing is a good choice for us.
    现在我们需要找到一种快速但足够强大的方法来检索语料库中与我们的问题最密切相关的文档。索引对我们来说是一个不错的选择。

    The first step is embed each document into a vector. We use bge-base-en-v1.5 as our embedding model.
    第一步是将每个文档嵌入到一个向量中。我们使用 bge-large-zh-v1.5 作为我们的embedding模型。
    """
    model = load_flag_model()
    embeddings = model.encode(corpus, convert_to_numpy=True)
    print(f"embeddings.shape: {embeddings.shape}")
    """
    Then, let's create a Faiss index and add all the vectors into it.
    然后，让我们创建一个 Faiss 索引并将所有向量添加到其中。

    If you want to know more about Faiss, refer to the tutorial of Faiss and indexing.
    如果您想了解更多关于 Faiss 的信息，请参考 Faiss 和索引教程。
    """
    index = faiss.IndexFlatIP(embeddings.shape[1])

    index.add(embeddings)
    print(f"index.ntotal: {index.ntotal}")

    # 3. Retrieve and Generate
    """
    Now we come to the most exciting part. Let's first embed our query and retrieve 3 most relevant document from it:
    现在我们来到最激动人心的部分。让我们首先嵌入我们的查询并从中检索 3 个最相关的文档：
    """
    q_embedding = model.encode_queries([user_input], convert_to_numpy=True)

    D, I = index.search(q_embedding, 3)
    res = np.array(corpus)[I]
    print(f"res: {res}")

    # Then set up the prompt for the chatbot:
    prompt="""
    You are a bot that makes recommendations for restaurants. 
    Please be brief, answer in short sentences without extra information.

    These are the restaurants list:
    {recommended_activities}

    The user's preference is: {user_input}
    Provide the user with 2 recommended restaurants based on the user's preference.
    """
    client,model_name = get_openai_client()

    response = client.chat.completions.create(
        model=model_name,
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": prompt.format(user_input=user_input, recommended_activities=res)
            }
        ]
    ).choices[0].message
    print(response.content)


